from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_deepseek import ChatDeepSeek
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.memory import ConversationBufferMemory
from langchain_community.agent_toolkits.load_tools import load_tools

load_dotenv(verbose=True)

llm = ChatDeepSeek(
    model="deepseek-chat",
    temperature=0.5
)

# tools = load_tools(["serpapi", "llm-math"], llm=llm)
tools = load_tools(["llm-math"], llm=llm)
memory = ConversationBufferMemory(
    memory_key="chat history",
    return_messages=True
)

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    memory=memory
)

print(agent.invoke("请问现任的中国主席是谁？他年龄的平方是多少？"))
print(agent.invoke("他的故乡是哪里？"))